#!/usr/bin/env python3
"""
AI医疗问诊助手 - Gradio前端界面
集成语音识别、对话分析、诊断建议和问诊记录生成功能
"""

import gradio as gr
import asyncio
import json
import tempfile
from pathlib import Path
from typing import Dict, List, Any
import sys

# 添加src目录到Python路径
sys.path.insert(0, str(Path(__file__).parent / "src"))

from config.settings import settings
from utils.logger import setup_logging
from asr.audio_processor import AudioProcessor
from nlp.dialogue_analyzer import DialogueAnalyzer
from services.record_service import RecordGenerator
from services.diagnosis_service import DiagnosisAdvisor

# 设置日志
setup_logging()

class MedicalAssistantApp:
    """医疗问诊助手应用类"""
    
    def __init__(self):
        self.audio_processor = AudioProcessor()
        self.dialogue_analyzer = DialogueAnalyzer()
        self.record_generator = RecordGenerator()
        self.diagnosis_advisor = DiagnosisAdvisor()
        self.current_session = {}
        
    async def initialize(self):
        """初始化所有组件"""
        print("🔄 初始化医疗问诊助手组件...")
        
        # 初始化语音处理器
        await self.audio_processor.initialize()
        print("✅ 语音处理器初始化完成")
        
        # 初始化对话分析器
        await self.dialogue_analyzer.initialize()
        print("✅ 对话分析器初始化完成")
        
        # 初始化诊断建议器
        await self.diagnosis_advisor.initialize()
        print("✅ 诊断建议器初始化完成")
        
        print("🎉 所有组件初始化完成")
    
    async def process_audio(self, audio_file: str) -> Dict[str, Any]:
        """处理音频文件并返回转录结果"""
        if not audio_file:
            return {"error": "请提供音频文件"}
        
        try:
            print(f"🎤 处理音频文件: {audio_file}")
            transcription = await self.audio_processor.transcribe_audio(audio_file)
            
            # 保存到当前会话
            self.current_session['transcription'] = transcription
            self.current_session['audio_file'] = audio_file
            
            return transcription
            
        except Exception as e:
            return {"error": f"音频处理失败: {str(e)}"}
    
    async def analyze_dialogue(self, dialogue_text: str) -> Dict[str, Any]:
        """分析医患对话文本"""
        if not dialogue_text:
            return {"error": "请输入对话文本"}
        
        try:
            print(f"🔍 分析对话文本: {dialogue_text[:50]}...")
            analysis_result = await self.dialogue_analyzer.analyze_dialogue(dialogue_text)
            
            # 保存到当前会话
            self.current_session['analysis'] = analysis_result
            self.current_session['dialogue_text'] = dialogue_text
            
            return analysis_result
            
        except Exception as e:
            return {"error": f"对话分析失败: {str(e)}"}
    
    async def generate_record(self) -> Dict[str, Any]:
        """生成问诊记录"""
        if 'analysis' not in self.current_session:
            return {"error": "请先进行对话分析"}
        
        try:
            print("📝 生成问诊记录...")
            record = await self.record_generator.generate_record(self.current_session['analysis'])
            
            # 保存到当前会话
            self.current_session['medical_record'] = record
            
            return record
            
        except Exception as e:
            return {"error": f"问诊记录生成失败: {str(e)}"}
    
    async def get_diagnosis(self) -> Dict[str, Any]:
        """获取诊断建议"""
        if 'medical_record' not in self.current_session:
            return {"error": "请先生成问诊记录"}
        
        try:
            print("🩺 获取诊断建议...")
            diagnosis = await self.diagnosis_advisor.get_diagnosis(self.current_session['medical_record'])
            
            # 保存到当前会话
            self.current_session['diagnosis'] = diagnosis
            
            return diagnosis
            
        except Exception as e:
            return {"error": f"诊断建议生成失败: {str(e)}"}
    
    def format_analysis_result(self, result: Dict[str, Any]) -> str:
        """格式化分析结果用于显示"""
        if 'error' in result:
            return f"❌ 错误: {result['error']}"
        
        output = []
        output.append("## 📊 对话分析结果")
        output.append("")
        
        # 说话人信息
        speakers = result.get('speakers', [])
        output.append(f"**👥 说话人**: {len(speakers)} 人")
        for speaker in speakers:
            output.append(f"  - {speaker.get('role', '未知')}: {speaker.get('name', '未知')}")
        
        # 医疗实体
        entities = result.get('medical_entities', [])
        output.append(f"**🏥 医疗实体**: {len(entities)} 个")
        
        # 按类型分组显示
        entity_types = {}
        for entity in entities:
            e_type = entity.get('type', 'other')
            if e_type not in entity_types:
                entity_types[e_type] = []
            entity_types[e_type].append(entity.get('text', ''))
        
        for e_type, items in entity_types.items():
            output.append(f"  **{e_type}**: {', '.join(items[:5])}{'...' if len(items) > 5 else ''}")
        
        # 严重程度
        severity = result.get('severity_score', 1)
        output.append(f"**⚠️ 严重程度**: {severity}/5")
        
        # 时间信息
        if 'timeline' in result:
            output.append(f"**⏱️ 时间线**: {len(result['timeline'])} 个事件")
        
        return "\n".join(output)
    
    def format_diagnosis_result(self, result: Dict[str, Any]) -> str:
        """格式化诊断结果用于显示"""
        if 'error' in result:
            return f"❌ 错误: {result['error']}"
        
        output = []
        output.append("## 🩺 诊断建议")
        output.append("")
        
        # 诊断建议
        suggestions = result.get('diagnosis_suggestions', [])
        output.append(f"**📋 诊断建议**: {len(suggestions)} 个")
        
        for i, suggestion in enumerate(suggestions, 1):
            disease = suggestion.get('disease', '未知疾病')
            probability = suggestion.get('probability', 0)
            evidence = suggestion.get('evidence', [])
            
            output.append(f"{i}. **{disease}** (概率: {probability:.2f})")
            if evidence:
                output.append(f"   证据: {', '.join(evidence[:3])}{'...' if len(evidence) > 3 else ''}")
        
        # 风险评估
        risk = result.get('risk_assessment', {})
        output.append(f"**⚠️ 风险评估**:")
        output.append(f"  紧急程度: {risk.get('urgency_level', '未知')}")
        output.append(f"  总体风险: {risk.get('overall_risk', '未知')}")
        
        # 建议措施
        recommendations = result.get('recommendations', [])
        if recommendations:
            output.append(f"**💡 建议措施**:")
            for rec in recommendations[:3]:
                output.append(f"  - {rec}")
        
        return "\n".join(output)
    
    def format_medical_record(self, record: Dict[str, Any]) -> str:
        """格式化问诊记录用于显示"""
        if 'error' in record:
            return f"❌ 错误: {record['error']}"
        
        output = []
        output.append("## 📝 问诊记录")
        output.append("")
        
        # 基本信息
        if 'basic_info' in record:
            output.append("### 基本信息")
            for key, value in record['basic_info'].items():
                output.append(f"- **{key}**: {value}")
            output.append("")
        
        # 主诉
        if 'chief_complaint' in record:
            output.append("### 主诉")
            output.append(record['chief_complaint'])
            output.append("")
        
        # 现病史
        if 'present_illness' in record:
            output.append("### 现病史")
            output.append(record['present_illness'])
            output.append("")
        
        # 体格检查
        if 'physical_exam' in record:
            output.append("### 体格检查")
            for item in record['physical_exam']:
                output.append(f"- {item}")
            output.append("")
        
        # 初步诊断
        if 'preliminary_diagnosis' in record:
            output.append("### 初步诊断")
            for diagnosis in record['preliminary_diagnosis']:
                output.append(f"- {diagnosis}")
        
        return "\n".join(output)

# 创建应用实例
app = MedicalAssistantApp()

async def gradio_process_audio(audio_file):
    """Gradio音频处理函数"""
    result = await app.process_audio(audio_file)
    if 'error' in result:
        return result['error'], "", "", ""
    
    transcription = result.get('text', '')
    return transcription, transcription, "", ""

async def gradio_analyze_dialogue(dialogue_text):
    """Gradio对话分析函数"""
    result = await app.analyze_dialogue(dialogue_text)
    formatted_result = app.format_analysis_result(result)
    return formatted_result

async def gradio_generate_record():
    """Gradio生成记录函数"""
    result = await app.generate_record()
    formatted_result = app.format_medical_record(result)
    return formatted_result

async def gradio_get_diagnosis():
    """Gradio获取诊断函数"""
    result = await app.get_diagnosis()
    formatted_result = app.format_diagnosis_result(result)
    return formatted_result

def create_gradio_interface():
    """创建Gradio界面"""
    with gr.Blocks(
        title="AI医疗问诊助手",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container { max-width: 1200px; margin: 0 auto; }
        .success { color: green; }
        .warning { color: orange; }
        .error { color: red; }
        """
    ) as demo:
        gr.Markdown("# 🏥 AI医疗问诊助手")
        gr.Markdown("基于语音识别和DeepSeek AI的智能医疗问诊辅助系统")
        
        with gr.Tab("🎤 语音输入"):
            with gr.Row():
                with gr.Column(scale=1):
                    audio_input = gr.Audio(
                        sources=["microphone", "upload"],
                        type="filepath",
                        label="录制或上传音频"
                    )
                    process_btn = gr.Button("🚀 处理音频", variant="primary")
                
                with gr.Column(scale=2):
                    transcription_output = gr.Textbox(
                        label="转录文本",
                        lines=6,
                        max_lines=12,
                        interactive=True
                    )
            
            process_btn.click(
                gradio_process_audio,
                inputs=[audio_input],
                outputs=[
                    transcription_output,
                    transcription_output,
                    gr.Textbox(visible=False),
                    gr.Textbox(visible=False)
                ]
            )
        
        with gr.Tab("🔍 对话分析"):
            with gr.Row():
                with gr.Column(scale=1):
                    dialogue_input = gr.Textbox(
                        label="医患对话文本",
                        placeholder="例如: 医生：您有什么不舒服？患者：我头痛三天了...",
                        lines=6,
                        max_lines=12
                    )
                    analyze_btn = gr.Button("📊 分析对话", variant="primary")
                
                with gr.Column(scale=2):
                    analysis_output = gr.Markdown(
                        label="分析结果"
                    )
            
            analyze_btn.click(
                gradio_analyze_dialogue,
                inputs=[dialogue_input],
                outputs=[analysis_output]
            )
        
        with gr.Tab("📝 问诊记录"):
            with gr.Row():
                record_btn = gr.Button("📋 生成问诊记录", variant="primary")
                record_output = gr.Markdown(
                    label="问诊记录"
                )
            
            record_btn.click(
                gradio_generate_record,
                inputs=[],
                outputs=[record_output]
            )
        
        with gr.Tab("🩺 诊断建议"):
            with gr.Row():
                diagnosis_btn = gr.Button("💡 获取诊断建议", variant="primary")
                diagnosis_output = gr.Markdown(
                    label="诊断建议"
                )
            
            diagnosis_btn.click(
                gradio_get_diagnosis,
                inputs=[],
                outputs=[diagnosis_output]
            )
        
        with gr.Tab("ℹ️ 关于"):
            gr.Markdown("""
            ## 关于AI医疗问诊助手
            
            **版本**: 1.0.0
            **技术栈**:
            - 🎤 语音识别: OpenAI Whisper
            - 🧠 AI分析: DeepSeek API
            - 🎨 前端: Gradio
            - 🐍 后端: Python + AsyncIO
            
            **功能特点**:
            - 实时语音转录和医学术语识别
            - 智能对话分析和医疗实体提取
            - 自动问诊记录生成
            - AI辅助诊断建议
            
            **注意事项**:
            - 本系统为辅助工具，不能替代专业医疗诊断
            - 请在实际医疗场景中谨慎使用
            - 所有诊断建议仅供参考
            """)
        
        # 添加状态显示
        gr.Markdown("---")
        gr.Markdown("### 📊 系统状态")
        status_text = gr.Textbox(
            label="运行状态",
            value="🟢 系统就绪",
            interactive=False
        )
    
    return demo

async def main():
    """主函数"""
    print("🚀 启动AI医疗问诊助手...")
    
    # 初始化应用
    await app.initialize()
    
    # 创建Gradio界面
    demo = create_gradio_interface()
    
    # 启动服务
    print("🌐 启动Gradio服务...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )

if __name__ == "__main__":
    asyncio.run(main())